"""代码生成智能体"""

from typing import Dict, List, Any, Set
from ..base.base_agent import BaseAgent
from ..base.agent_state import Task, TaskStatus
from ...llm.llm_manager import llm_manager


class CodeGeneratorAgent(BaseAgent):
    """代码生成智能体 - 负责根据需求生成代码"""
    
    def __init__(self, agent_id: str = "code_generator"):
        capabilities = {
            "code_generation",
            "syntax_analysis",
            "pattern_recognition",
            "language_support",
            "code_optimization"
        }
        super().__init__(agent_id, "代码生成智能体", capabilities)
        self.supported_languages = {
            "python", "javascript", "typescript", "java", 
            "cpp", "csharp", "go", "rust", "php"
        }
    
    async def _setup(self) -> None:
        """设置代码生成器"""
        self.logger.info("代码生成智能体初始化完成")
    
    async def _execute_task(self, task: Task) -> Dict[str, Any]:
        """执行代码生成任务"""
        task_type = task.metadata.get("type", "unknown")
        
        if task_type == "generate_code":
            return await self._generate_code(task)
        elif task_type == "generate_tests":
            return await self._generate_tests(task)
        elif task_type == "generate_docs":
            return await self._generate_docs(task)
        else:
            raise ValueError(f"未知的任务类型: {task_type}")
    
    async def _generate_code(self, task: Task) -> Dict[str, Any]:
        """生成代码"""
        requirements = task.metadata.get("requirements", {})
        language = requirements.get("language", "python")
        functionality = requirements.get("functionality", "")
        
        if language not in self.supported_languages:
            raise ValueError(f"不支持的语言: {language}")
        
        # 构建提示词
        prompt = f"""
        请根据以下需求生成{language}代码：
        
        功能描述：{functionality}
        编程语言：{language}
        
        要求：
        1. 代码要完整可运行
        2. 包含必要的注释
        3. 遵循{language}的最佳实践
        4. 包含错误处理
        
        请直接返回代码，不要包含其他说明文字。
        """
        
        messages = [
            {"role": "system", "content": "你是一个专业的代码生成助手，能够根据需求生成高质量的代码。"},
            {"role": "user", "content": prompt}
        ]
        
        try:
            # 使用百炼平台生成代码
            response = llm_manager.chat_completion("bailian", messages, model="qwen-plus")
            
            generated_code = response.choices[0].message.content
            
            return {
                "language": language,
                "code": generated_code,
                "requirements": requirements,
                "status": "success"
            }
            
        except Exception as e:
            self.logger.error(f"代码生成失败: {str(e)}")
            return {
                "language": language,
                "code": "",
                "error": str(e),
                "status": "failed"
            }
    
    async def _generate_tests(self, task: Task) -> Dict[str, Any]:
        """生成测试代码"""
        source_code = task.metadata.get("source_code", "")
        test_framework = task.metadata.get("test_framework", "pytest")
        
        prompt = f"""
        请为以下代码生成{test_framework}测试用例：
        
        ```python
        {source_code}
        ```
        
        要求：
        1. 测试覆盖主要功能
        2. 包含边界情况测试
        3. 使用{test_framework}框架
        4. 测试用例要完整可运行
        """
        
        messages = [
            {"role": "system", "content": "你是一个专业的测试代码生成助手。"},
            {"role": "user", "content": prompt}
        ]
        
        try:
            response = llm_manager.chat_completion("bailian", messages, model="qwen-plus")
            test_code = response.choices[0].message.content
            
            return {
                "test_framework": test_framework,
                "test_code": test_code,
                "source_code": source_code,
                "status": "success"
            }
            
        except Exception as e:
            self.logger.error(f"测试代码生成失败: {str(e)}")
            return {
                "test_framework": test_framework,
                "test_code": "",
                "error": str(e),
                "status": "failed"
            }
    
    async def _generate_docs(self, task: Task) -> Dict[str, Any]:
        """生成文档"""
        source_code = task.metadata.get("source_code", "")
        doc_format = task.metadata.get("format", "markdown")
        
        prompt = f"""
        请为以下代码生成{doc_format}格式的文档：
        
        ```python
        {source_code}
        ```
        
        要求：
        1. 包含函数/类的说明
        2. 包含参数说明
        3. 包含返回值说明
        4. 包含使用示例
        5. 使用{doc_format}格式
        """
        
        messages = [
            {"role": "system", "content": "你是一个专业的代码文档生成助手。"},
            {"role": "user", "content": prompt}
        ]
        
        try:
            response = llm_manager.chat_completion("bailian", messages, model="qwen-plus")
            documentation = response.choices[0].message.content
            
            return {
                "format": doc_format,
                "documentation": documentation,
                "source_code": source_code,
                "status": "success"
            }
            
        except Exception as e:
            self.logger.error(f"文档生成失败: {str(e)}")
            return {
                "format": doc_format,
                "documentation": "",
                "error": str(e),
                "status": "failed"
            }
